🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1458

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 8a9a644 with merge base e5cf6e5 ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

cc: @manuelcandales

Can have please for MPS? 🥺🥺 (Separate PR)

Awaiting pytorch/executorch#7759

@Jack-Khuu did you update the version AO uses in ET?

Yup pytorch/executorch@9836b39 Points to pytorch/ao@11333ba

Hello @metascroy @Jack-Khuu , what is the plan to get this in mainline? We would like to use KleidiAI kernels from aten via this quantizer path. Let us know if we need to raise a new PR ?

Hi @nikhil-arm, we're still planning to land this

Can you share the specific commit hashes y'all need?

@nikhil-arm We've bumped the AO pin on main.
Please let me know if you there's any additional support needed to unblock KleidiAI kernels

After a suite of rebases, pinbumps, and splitting up tests up we know what we're tackling:

• test-torchao-experimental-cpp (macos-14-xlarge): Tests the AOTI runner and likely failing (also in main) due to not linking to the LibOMP from torch as @malfet mentioned in Bump PT 2025131 and ET pins 20250209 #1493.

• test-torchao-experimental-et (macos-14-xlarge): Tests the ET runner; looks like a install bug where USE_CPP isn't set, but will likely run into the same LibOMP issue as above

Hello @metascroy @Jack-Khuu , what is the plan to get this in mainline? We would like to use KleidiAI kernels from aten via this quantizer path. Let us know if we need to raise a new PR ?

Sorry about the delay @nikhil-arm.

@Jack-Khuu let's try to get this landed within the next week. Bumping the ao pin in torchchat had various conflicts with the CI, but I think we can dedicate to making this work.

I think it does make sense to first land pytorch/ao#1836 in torchao before bumping because they've already deprecated the old quantizers in quantize_.

Hi @nikhil-arm The experimental kernels from AO are now in. Can you share a link to the KleidiAI Kernels (either here or a GH issue if you can share more context)?

Hi @nikhil-arm The experimental kernels from AO are now in. Can you share a link to the KleidiAI Kernels (either here or a GH issue if you can share more context)?

For context for both @nikhil-arm and @Jack-Khuu. There are two locations of KleidiAI kernels in PyTorch/torchao now.

PyTorch has a 4-bit quantized linear op backed by KleidiAI. I think @nikhil-arm wants to enable this in torchchat. To do this, you need to pass "aten" to the target in PackedLinearInt8DynamicActivationIntxWeightLayout(target="aten") here: https://github.com/pytorch/torchchat/blob/main/torchchat/utils/quantize.py#L150. I will leave it to @nikhil-arm to put up a PR to pipe this through. This should then work with most surfaces, except ExecuTorch, although I don't know if it's been tested in anything other than eager mode.

We have also pulled in KleidiAI kernels in torchao itself. These should work on all surfaces, including ExecuTorch. To enable these, we need to add the flag TORCHAO_BUILD_KLEIDIAI=1 before the pip install:

When enabled, PackedLinearInt8DynamicActivationIntxWeightLayout() will use the KleidiAI kernel when supported (4-quantization, has_weight_zeros=false), and use our native torchao kernels in other cases. You can see what kernel is being selected at runtime by setting the environment variable TORCH_CPP_LOG_LEVEL=Info. Eventually we want to enable this flag by default, but I have it disabled right now for two reasons: 1) the flag increases the install time because it clones/builds KleidiAI; 2) there is limited perf benefit to enabling it now based on the KleidiAI kernels we have enabled in torchao (GEMV neondot kernels) vs. the native torchao Arm kernels (also GEMV neondot kernels).